import torch
import pandas as pd
from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer
from pytorch_forecasting.metrics import RMSE
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import EarlyStopping

# 读取数据（假设数据文件名为 sales_data.csv）
df = pd.read_csv('lilylikes-all-LO2312019AXNHS.csv', low_memory=False, encoding='GBK')

# 日期转换与特征工程
df["date_id"] = pd.to_datetime(df["date_id"], format="%Y%m%d")
df["first_new_date"] = pd.to_datetime(df["first_new_date"], format="%Y/%m/%d")
df["product_age"] = (df["date_id"] - df["first_new_date"]).dt.days

# 添加时间相关特征
df["month"] = df.date_id.dt.month.astype(str).astype("category")
df["day_of_week"] = df.date_id.dt.dayofweek.astype(str).astype("category")
df["day"] = df.date_id.dt.day.astype(str).astype("category")

# 创建时间索引列
df["time_idx"] = (df["date_id"] - df["date_id"].min()).dt.days

# 按产品分组排序
df = df.sort_values(by=["product_code", "date_id"]).reset_index(drop=True)

# 构造滞后特征（过去3天销量）
df["lag_3"] = df.groupby("product_code")["sale_qty"].shift(3)

# 过滤有效数据
df = df[df["product_age"] >= 0].dropna()


def fill_missing_dates(df):

    # 生成完整日期范围
    min_date = df['date_id'].min()
    max_date = df['date_id'].max()
    all_dates = pd.date_range(min_date, max_date, freq='D')

    # 获取所有 product_code
    # product_codes = ""

    # 创建完整日期和 product_code 的组合
    full_df = pd.DataFrame(all_dates, columns=['date_id'])

    # 合并原始数据
    full_df = full_df.merge(df, on=['date_id'], how='left')

    # 填充缺失的 sale_qty 为 0
    full_df['sale_qty'] = full_df['sale_qty'].fillna(0)

    return full_df

fill_missing_dates(df)

max_prediction_length = 7  # 预测未来7天
max_encoder_length = 21    # 使用过去21天数据

training = TimeSeriesDataSet(
    df,
    time_idx="time_idx",
    target="sale_qty",
    group_ids=["product_code"],  # 每个产品的独立序列
    static_categoricals=["middle_class_name"],  # 静态分类特征
    time_varying_known_categoricals=["month", "day_of_week"],  # 已知时间特征
    time_varying_known_reals=["product_age", "lag_3"],  # 已知连续特征
    time_varying_unknown_reals=[],  # 未知动态特征（这里为空）
    max_encoder_length=max_encoder_length,
    max_prediction_length=max_prediction_length,
    target_normalizer=None,  # 保持原始销量数值
    add_relative_time_idx=True,
    add_target_scales=False,
    allow_missing_timesteps=True  # 允许缺失的时间步长
)

# 创建数据加载器
batch_size = 64
train_dataloader = training.to_dataloader(
    train=True,
    batch_size=batch_size,
    num_workers=2
)



tft = TemporalFusionTransformer.from_dataset(
    training,
    hidden_size=32,  # 调整模型容量
    lstm_layers=2,   # 增加LSTM层数
    attention_head_size=4,
    dropout=0.1,
    learning_rate=0.03,
    output_size=1,
    loss=RMSE()
)



# 配置早停机制
early_stop_callback = EarlyStopping(
    monitor="val_loss",
    min_delta=1e-4,
    patience=10,
    verbose=False,
    mode="min"
)

# 训练配置
trainer = Trainer(
    max_epochs=50,
    accelerator="auto",
    enable_model_summary=True,
    callbacks=[early_stop_callback],
    limit_train_batches=50  # 每epoch使用50个batch
)

# 拆分验证集
validation = TimeSeriesDataSet.from_dataset(
    training,
    df,
    predict=True,
    stop_randomization=True
)
val_dataloader = validation.to_dataloader(
    train=False,
    batch_size=batch_size,
    num_workers=2
)

# 开始训练
trainer.fit(
    tft,
    train_dataloaders=train_dataloader,
    val_dataloaders=val_dataloader
)


# 生成预测
raw_predictions, _ = tft.predict(
    val_dataloader,
    mode="raw",
    return_x=True
)

# 可视化样例
for idx in range(3):  # 展示前3个产品的预测
    tft.plot_prediction(
        raw_predictions.x,
        raw_predictions.output,
        idx=idx,
        add_loss_to_title=True
    )
